from ship_envs.core.obs_viewer import ObsViewer
from ship_envs.core.ship import Ship
from ship_envs.environment.imazu_problem_env_factory import new_case_env_by_index
import os

def env_test(env):
    env.reset()
    while True:
        obs, reward, done, info = env.step(env.action_space.sample())
        env.render()
        if done:
            break

def show_env_obs(env, save=None):
    env.reset()
    for _ in range(1):
        env.step(0)
    viewer = ObsViewer(500, 500)
    entitys = env.world.get_entitys()
    entitys.reverse()
    for entity in entitys:
        viewer.add_model(entity.model)
        if isinstance(entity, Ship):
            viewer.add_model(entity.target)
    img = viewer.get_img()
    if save:
        dir_path = r"C:\Users\yilun\whut\毕业相关\实验环境"
        img_path = os.path.join(dir_path, str(save) + '.jpg')
        print(img_path)
        img.save(img_path)
    else:
        img.show()

if __name__ == '__main__':
    case = None
    for i in range(1, 23):
        # env_test(new_case_env_by_index(i))
        show_env_obs(new_case_env_by_index(i), save=i)